# coding=utf-8
# encoding=utf-8

import json

from kafka import KafkaProducer
from kafka import KafkaConsumer
from kafka.errors import KafkaError


class PicKafkaProducer(object):
    """
     生产模块：根据不同的key，区分消息
    """
    def __init__(self, kafka_host, kafka_port, kafka_topic, key):
        self.kafka_host = kafka_host
        self.kafka_port = kafka_port
        self.kafka_topic = kafka_topic
        self.key = key
        self.producer = KafkaProducer(bootstrap_servers='{kafka_host}:{kafka_port}'.format(
            kafka_host=self.kafka_host,
            kafka_port=self.kafka_port)
        )

    def send_json_data(self, params):
        try:
            parmas_message = json.dumps(params)
            producer = self.producer
            producer.send(self.kafka_topic, key=self.key, value=parmas_message.encode('utf-8'))
            producer.flush()
        except KafkaError as e:
            print(e)


class PicKafkaConsumer(object):
    """
    消费模块: 通过不同groupid消费topic里面的消息
    """

    def __init__(self, kafka_host, kafka_port, kafka_topic, group_id):
        self.kafka_host = kafka_host
        self.kafka_port = kafka_port
        self.kafka_topic = kafka_topic
        self.group_id = group_id
        self.consumer = KafkaConsumer(self.kafka_topic, group_id=self.group_id,
                                      bootstrap_servers='{kafka_host}:{kafka_port}'.format(
                                          kafka_host=self.kafka_host,
                                          kafka_port=self.kafka_port)
                                      )

    def consume_data(self):
        try:
            for message in self.consumer:
                yield message
        except KeyboardInterrupt as e:
            print(e)
